#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <openssl/evp.h>
#include <openssl/pem.h>
#include <openssl/rand.h>
#include <openssl/err.h>
#include <openssl/ec.h>

// 默认 SM2 用户 ID
const unsigned char SM2_DEFAULT_USERID[] = "1234567812345678";
const size_t SM2_DEFAULT_USERID_LEN = 16;

// 错误处理函数
void handleErrors(void) {
    fprintf(stderr, "OpenSSL Error:\n");
    ERR_print_errors_fp(stderr);
    abort();
}

// 从 PEM 文件加载公钥或私钥
EVP_PKEY *load_pkey(const char *filename, int is_private) {
    BIO *bio = BIO_new_file(filename, "r");
    if (!bio) {
        perror("BIO_new_file");
        return NULL;
    }

    EVP_PKEY *pkey = NULL;
    if (is_private) {
        pkey = PEM_read_bio_PrivateKey(bio, NULL, NULL, NULL);
    } else {
        pkey = PEM_read_bio_PUBKEY(bio, NULL, NULL, NULL);
    }

    BIO_free(bio);
    if (!pkey) {
        fprintf(stderr, "Error loading key from %s\n", filename);
        handleErrors();
    }
    return pkey;
}

// 将数据写入文件
void write_to_file(const char *filename, const unsigned char *data, size_t len) {
    FILE *fp = fopen(filename, "wb"); // "wb" - write binary
    if (!fp) {
        perror(filename);
        return;
    }
    if (fwrite(data, 1, len, fp) != len) {
        fprintf(stderr, "Error writing to %s\n", filename);
    }
    fclose(fp);
    printf("Alice: Wrote %zu bytes to %s\n", len, filename);
}

/**
 * @brief Alice 的发送操作 (函数原型与之前相同)
 */
int alice_send(EVP_PKEY *Pkb, EVP_PKEY *Ska, const unsigned char *P, size_t P_len,
               unsigned char **C, size_t *C_len,
               unsigned char **KC, size_t *KC_len,
               unsigned char **S1, size_t *S1_len) {

    printf("--- Alice 开始发送 ---\n");

    // 1. sm4 key 使用...产生, 16字节, 记作k
    unsigned char k[16];
    if (RAND_bytes(k, sizeof(k)) != 1) {
        fprintf(stderr, "Alice: 产生SM4密钥k失败\n");
        handleErrors();
        return 0;
    }
    printf("Alice: 成功生成 16 字节 SM4 密钥 k\n");

    // 2. Sm4Enc(k, P) = C (使用 SM4-CBC)
    unsigned char iv[16];
    if (RAND_bytes(iv, sizeof(iv)) != 1) {
        fprintf(stderr, "Alice: 产生SM4 IV失败\n");
        handleErrors();
        return 0;
    }

    EVP_CIPHER_CTX *ctx_sm4 = EVP_CIPHER_CTX_new();
    if (!ctx_sm4) handleErrors();

    *C = malloc(16 + P_len + 16); // IV + Data + Padding
    if (!*C) {
        perror("malloc C");
        return 0;
    }
    
    memcpy(*C, iv, 16); // 预置 IV
    unsigned char *C_data = *C + 16;
    int out_len = 0, final_len = 0;

    if (EVP_EncryptInit_ex(ctx_sm4, EVP_sm4_cbc(), NULL, k, iv) != 1) handleErrors();
    if (EVP_EncryptUpdate(ctx_sm4, C_data, &out_len, P, P_len) != 1) handleErrors();
    if (EVP_EncryptFinal_ex(ctx_sm4, C_data + out_len, &final_len) != 1) handleErrors();

    *C_len = 16 + out_len + final_len; // 总长度 = IV + 密文
    EVP_CIPHER_CTX_free(ctx_sm4);
    printf("Alice: Sm4Enc(k, P) = C 完成 (C 包含 IV), C 总长度: %zu 字节\n", *C_len);

    // 3. Sm2Enc(Pkb, k) = KC (SM2 公钥加密)
    EVP_PKEY_CTX *ctx_sm2_enc = EVP_PKEY_CTX_new_from_pkey(NULL, Pkb, NULL);
    if (!ctx_sm2_enc) handleErrors();
    if (EVP_PKEY_encrypt_init(ctx_sm2_enc) <= 0) handleErrors();

    if (EVP_PKEY_encrypt(ctx_sm2_enc, NULL, KC_len, k, sizeof(k)) <= 0) handleErrors();
    *KC = malloc(*KC_len);
    if (!*KC) {
        perror("malloc KC");
        return 0;
    }
    if (EVP_PKEY_encrypt(ctx_sm2_enc, *KC, KC_len, k, sizeof(k)) <= 0) handleErrors();
    
    printf("Alice: Sm2Enc(Pkb, k) = KC 完成, KC 长度: %zu 字节\n", *KC_len);
    EVP_PKEY_CTX_free(ctx_sm2_enc);

    // 4. Sm2Sign(Ska, C) = S1 (SM2 签名, 使用 SM3)
    EVP_MD_CTX *ctx_md = EVP_MD_CTX_new();
    EVP_PKEY_CTX *pctx_sig = NULL;
    if (!ctx_md) handleErrors();

    if (EVP_DigestSignInit(ctx_md, &pctx_sig, EVP_sm3(), NULL, Ska) <= 0) handleErrors();
    if (EVP_PKEY_CTX_set1_id(pctx_sig, SM2_DEFAULT_USERID, SM2_DEFAULT_USERID_LEN) <= 0) handleErrors();
    if (EVP_DigestSignUpdate(ctx_md, *C, *C_len) <= 0) handleErrors();

    if (EVP_DigestSignFinal(ctx_md, NULL, S1_len) <= 0) handleErrors();
    *S1 = malloc(*S1_len);
    if (!*S1) {
        perror("malloc S1");
        return 0;
    }
    if (EVP_DigestSignFinal(ctx_md, *S1, S1_len) <= 0) handleErrors();

    printf("Alice: Sm2Sign(Ska, C) = S1 完成, S1 长度: %zu 字节\n", *S1_len);
    EVP_MD_CTX_free(ctx_md);

    printf("--- Alice 步骤完成 ---\n");
    return 1;
}

int main() {
    // 请替换为您自己的8位学号+姓名
    const char *plaintext = "zsy-20231324";
    
    // 写入 plain.txt
    FILE *fp_plain = fopen("plain.txt", "w");
    if (!fp_plain) {
        perror("fopen plain.txt");
        return 1;
    }
    fputs(plaintext, fp_plain);
    fclose(fp_plain);
    printf("Alice: 已创建 plain.txt, 内容: '%s'\n\n", plaintext);

    // --- 密钥加载 (Alice 只需要 Bob 的公钥 和 自己的私钥) ---
    EVP_PKEY *skA = load_pkey("skA.pem", 1);
    EVP_PKEY *pkB = load_pkey("pkB.pem", 0);

    if (!skA || !pkB) {
        fprintf(stderr, "Alice: 加载密钥失败 (需要 skA.pem, pkB.pem)\n");
        return 1;
    }

    // --- 模拟传输 ---
    unsigned char *C = NULL;
    size_t C_len = 0;
    unsigned char *KC = NULL;
    size_t KC_len = 0;
    unsigned char *S1 = NULL;
    size_t S1_len = 0;

    // 1. Alice 执行发送操作
    if (!alice_send(pkB, skA, (const unsigned char*)plaintext, strlen(plaintext),
                    &C, &C_len, &KC, &KC_len, &S1, &S1_len)) {
        fprintf(stderr, "Alice 发送过程失败。\n");
    } else {
        // 2. 将结果写入文件
        write_to_file("envelope.C", C, C_len);
        write_to_file("envelope.KC", KC, KC_len);
        write_to_file("envelope.S1", S1, S1_len);
        printf("\nAlice: 数字信封已生成 (envelope.C, envelope.KC, envelope.S1)\n");
    }

    // --- 清理 ---
    EVP_PKEY_free(skA);
    EVP_PKEY_free(pkB);
    free(C);
    free(KC);
    free(S1);
    
    return 0;
}
